package ai.neuralnet;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;

public class FeatureSelector
{
    private ClassificationDataSet dataSet;

    private int categoryCount;

    private int rawFeatureCount;

    private double acceptanceThreshold;

    public FeatureSelector(ClassificationDataSet dataSet)
    {
        this.dataSet = dataSet;
        categoryCount = dataSet.getCategoryCount();
        rawFeatureCount = dataSet.getPointDimension();
    }

    public ClassificationDataSet selectFeatures()
    {
        boolean[] allCategories = new boolean[categoryCount];
        Arrays.fill(allCategories, true);
        return selectFeatures(allCategories);
    }

    public ClassificationDataSet selectFeatures(boolean[] categories)
    {
        List<ClassificationDataSet> categoryGroups = splitByCategory(dataSet, categoryCount);
        int finalCatCount = categoryCount;
        for (int i = categoryCount - 1; i >= 0; i--)
        {
            if (!categories[i])
            {
                categoryGroups.remove(i);
                finalCatCount--;
            }

        }

        List<double[]> categoryAvg = new ArrayList<double[]>();
        List<double[]> categoryStdDev = new ArrayList<double[]>();

        for (ClassificationDataSet groupSet : categoryGroups)
        {
            double[] avgPoint = new double[groupSet.getPointDimension()];
            double[] stdDev = new double[groupSet.getPointDimension()];

            for (double[] dataPoint : groupSet.getInputValues())
            {
                for (int dim = 0; dim < dataPoint.length; dim++)
                {
                    avgPoint[dim] += dataPoint[dim];
                }
            }

            for (int dim = 0; dim < avgPoint.length; dim++)
            {
                avgPoint[dim] /= groupSet.size();
            }

            for (double[] dataPoint : groupSet.getInputValues())
            {
                for (int dim = 0; dim < dataPoint.length; dim++)
                {
                    stdDev[dim] += Math.pow(dataPoint[dim] - avgPoint[dim], 2);
                }
            }

            for (int dim = 0; dim < stdDev.length; dim++)
            {
                stdDev[dim] = Math.sqrt(stdDev[dim] / groupSet.size());
            }

            categoryAvg.add(avgPoint);
            categoryStdDev.add(stdDev);
        }


        double[] featureQuality = getFeatureQuality(categoryAvg, categoryStdDev);
        int finalFeatureCount = 0;
        for (double d : featureQuality)
        {
            if (d > acceptanceThreshold) finalFeatureCount++;
        }


        ClassificationDataSet selectedSet = new ClassificationDataSet();

        for (int i = 0; i < dataSet.size(); i++)
        {
            double[] rawInput = dataSet.getInputValues().get(i);
            double[] outputs = dataSet.getOutputValues().get(i);

            double[] finalInput = new double[finalFeatureCount];
            int fI = 0;
            for (int k = 0; k < rawFeatureCount; k++)
                if (featureQuality[i] > acceptanceThreshold)
                {
                    finalInput[fI++] = rawInput[k];
                }

            selectedSet.addInputOutputPair(finalInput, outputs);
        }

        return selectedSet;
    }


    private double[] getFeatureQuality(List<double[]> categoryAvg, List<double[]> categoryStdDev)
    {
        double[] featureQuality = new double[rawFeatureCount];

        int finalCatCount = categoryAvg.size();

        int totalComparisons = (finalCatCount - 1) * finalCatCount / 2;

        for (int i = 0; i < finalCatCount; i++)
        {
            for (int j = 1; j < finalCatCount; j++)
            {
                double[] iAvg = categoryAvg.get(i);
                double[] jAvg = categoryAvg.get(j);
                double[] iStdDev = categoryStdDev.get(i);
                double[] jStdDev = categoryStdDev.get(j);

                for (int feature = 0; feature < rawFeatureCount; feature++)
                {
                    featureQuality[feature] += qualityScore(iAvg[feature], iStdDev[feature], jAvg[feature], jStdDev[feature]) / totalComparisons;
                }
            }
        }

        return featureQuality;
    }

    private static double qualityScore(double d1Center, double d1Err, double d2Center, double d2Err)
    {
        double lC = Math.min(d1Center, d2Center);
        double rC = Math.max(d1Center, d2Center);
        double lE = lC == d1Center ? d1Err : d2Err;
        double rE = lE == d1Err ? d2Err : d1Err;

        return ((rC - rE) - (lC + lE)) / ((rC + rE) - (lC - lE));

    }


    public static List<ClassificationDataSet> splitByCategory(ClassificationDataSet dataSet, int categoryCount)
    {
        List<ClassificationDataSet> categoryGroups = new ArrayList<ClassificationDataSet>();

        for (int i = 0; i < categoryCount; i++)
        {
            categoryGroups.add(new ClassificationDataSet());
        }

        for (int i = 0; i < dataSet.size(); i++)
        {
            int cat = dataSet.getCategory(i);
            categoryGroups.get(cat).addInputOutputPair(dataSet.getInputValues().get(i), dataSet.getOutputValues().get(i));
        }

        return categoryGroups;
    }

    public void setAcceptanceThreshold(double acceptanceThreshold)
    {
        this.acceptanceThreshold = acceptanceThreshold;
    }
}
